import os
import torch
import segmentation_models_pytorch as smp


def get_model(model_path=None):
    # if model_path == None:
    #     checkpoint = '/home/slz/project/P1/user_data/model_data/efficientnet-b4-6ed6700e.pth'
    #     if not os.path.exists(checkpoint):
    #         checkpoint = 'imagenet'
    # ['imagenet', 'advprop']
    encoder_name = 'efficientnet-b4'
    encoder_name = 'resnet50'
    model = smp.UnetPlusPlus(
        encoder_name=encoder_name,  # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
        encoder_weights='imagenet',  # use `imagenet` pretreined weights for encoder initialization
        in_channels=3,  # model input channels (1 for grayscale images, 3 for RGB, etc.)
        classes=1,  # model output channels (number of classes in your dataset)
    )
    # checkpoint = '/home/slz/.cache/torch/hub/checkpoints/efficientnet-b4-6ed6700e.pth'
    # model.encoder.load_state_dict(torch.load(checkpoint))
    # next(model.parameters()).device
    return model


if __name__ == "__main__":
    model = get_model()
    # print(model.encoder)
    # checkpoint = './round1/r1fold3_uppmodel_new3.pth'
    # model.load_state_dict(torch.load(checkpoint))